import queue
import threading

from django.db.backends.mysql.base import DatabaseWrapper as MysqlDatabaseWrapper

DEFAULT_DB_POOL_SIZE = 5


class ConnectPool:
    def __init__(self, database, conn_params, pool_size):
        self.conn_params = conn_params
        self.pool_size = pool_size
        self.count = 0
        self.connects = queue.Queue()
        self.database = database

    def get_connection(self):
        if self.connects.empty() and self.count < self.pool_size:
            self.count = self.count + 1
            conn = self.database.connect(**self.conn_params)
        else:
            conn = self.connects.get()
        # 检测连接是否有效，去掉性能更好，但建议保留
        conn.ping(reconnect=True)
        return conn

    def release_connection(self, conn):
        self.connects.put(conn)


class DatabaseWrapper(MysqlDatabaseWrapper):
    # """
    # 使用此库时绝对不能设置 CONN_MAX_AGE 连接参数，否则会造成使用连接后不会快速释放到连接池，从而造成连接池阻塞
    # """
    pool_size = None
    connect_pools = {}
    mutex = threading.Lock()

    def connect(self):
        """
        使用时绝对不能设置 CONN_MAX_AGE 连接参数，否则会造成使用连接后不会快速释放到连接池，从而造成连接池阻塞
        :return:
        """
        self.settings_dict['CONN_MAX_AGE'] = 0
        super(DatabaseWrapper, self).connect()

    def get_new_connection(self, conn_params):
        with self.mutex:
            # 获取DATABASES配置字典中的DB_POOL_SIZE参数
            if not self.pool_size:
                self.pool_size = self.settings_dict.get('DB_POOL_SIZE') or DEFAULT_DB_POOL_SIZE

            # 获取连接池
            if self.alias in self.connect_pools:
                connect_pool = self.connect_pools[self.alias]
            else:
                connect_pool = ConnectPool(self.Database, conn_params, self.pool_size)
                # 保存连接池实例
                self.connect_pools[self.alias] = connect_pool

            return connect_pool.get_connection()

    def _close(self):
        with self.mutex:
            # 覆盖掉原来的 close 方法，查询结束后连接释放回连接池
            if self.connection is not None:
                with self.wrap_database_errors:
                    # 获取连接池
                    connect_pool = self.connect_pools[self.alias]
                    # 释放连接到连接池
                    return connect_pool.release_connection(self.connection)
